﻿/*
 * KdTree.cs
 * RVO2 Library C#
 *
 * Copyright 2008 University of North Carolina at Chapel Hill
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
 * Please send all bug reports to <geom@cs.unc.edu>.
 *
 * The authors may be contacted via:
 *
 * Jur van den Berg, Stephen J. Guy, Jamie Snape, Ming C. Lin, Dinesh Manocha
 * Dept. of Computer Science
 * 201 S. Columbia St.
 * Frederick P. Brooks, Jr. Computer Science Bldg.
 * Chapel Hill, N.C. 27599-3175
 * United States of America
 *
 * <http://gamma.cs.unc.edu/RVO2/>
 */

using System;
using System.Collections.Generic;
using UnityEngine;

namespace Core.RVO
{
    /**
     * <summary>Defines k-D trees for agents and static obstacles.</summary>
     */
    public partial class KdTree : Component, IAwake
    {
        /**
         * <summary>The maximum size of an agent k-D tree leaf.</summary>
         */
        private const int MAX_LEAF_SIZE = 10;

        private readonly List<Agent> agents = new();
        private AgentTreeNode[] agentTree;
        private ObstacleTreeNode obstacleTree;

        private RVOSystem sys;

        public void Awake()
        {
            sys = Entity as RVOSystem;
        }

        /**
         * <summary>Builds an agent k-D tree.</summary>
         */
        public void BuildAgentTree()
        {
            int agentCnt = sys.Agents.Count;
            if (agents.Count != agentCnt)
            {
                agents.Clear();

                for (int i = 0; i < agentCnt; ++i)
                {
                    agents.Add(sys.Agents.GetByIndex(i));
                }

                agentTree = new AgentTreeNode[2 * agentCnt];

                for (int i = 0; i < agentTree.Length; ++i)
                {
                    agentTree[i] = new AgentTreeNode();
                }
            }

            if (agents.Count != 0)
            {
                BuildAgentTreeRecursive(0, agents.Count, 0);
            }
        }

        /**
         * <summary>Builds an obstacle k-D tree.</summary>
         */
        public void BuildObstacleTree()
        {
            obstacleTree = new ObstacleTreeNode();

            IList<Obstacle> obstacles = new List<Obstacle>(sys.Obstacles.Count);

            for (int i = 0; i < sys.Obstacles.Count; ++i)
            {
                obstacles.Add(sys.Obstacles.Get(i));
            }

            obstacleTree = BuildObstacleTreeRecursive(obstacles);
        }

        /**
         * <summary>Computes the agent neighbors of the specified agent.
         * </summary>
         *
         * <param name="agent">The agent for which agent neighbors are to be
         * computed.</param>
         * <param name="range_sq">The squared range around the agent.</param>
         */
        public void ComputeAgentNeighbors(Agent agent, ref float range_sq)
        {
            QueryAgentTreeRecursive(agent, ref range_sq, 0);
        }

        /**
         * <summary>Computes the obstacle neighbors of the specified agent.
         * </summary>
         *
         * <param name="agent">The agent for which obstacle neighbors are to be
         * computed.</param>
         * <param name="range_sq">The squared range around the agent.</param>
         */
        public void ComputeObstacleNeighbors(Agent agent, float range_sq)
        {
            QueryObstacleTreeRecursive(agent, range_sq, obstacleTree);
        }

        /**
         * <summary>Queries the visibility between two points within a specified
         * radius.</summary>
         *
         * <returns>True if q1 and q2 are mutually visible within the radius;
         * false otherwise.</returns>
         *
         * <param name="q1">The first point between which visibility is to be
         * tested.</param>
         * <param name="q2">The second point between which visibility is to be
         * tested.</param>
         * <param name="radius">The radius within which visibility is to be
         * tested.</param>
         */
        public bool QueryVisibility(Vector2 q1, Vector2 q2, float radius)
        {
            return QueryVisibilityRecursive(q1, q2, radius, obstacleTree);
        }

        public int QueryNearAgent(Vector2 point, float radius)
        {
            float rangeSq = float.MaxValue;
            int agentNo = -1;
            QueryAgentTreeRecursive(point, ref rangeSq, ref agentNo, 0);
            if (rangeSq < radius * radius)
                return agentNo;
            return -1;
        }

        /**
         * <summary>Recursive method for building an agent k-D tree.</summary>
         *
         * <param name="begin">The beginning agent k-D tree node index.</param>
         * <param name="end">The ending agent k-D tree node index.</param>
         * <param name="node">The current agent k-D tree node index.</param>
         */
        public void BuildAgentTreeRecursive(int begin, int end, int node)
        {
            agentTree[node].Begin = begin;
            agentTree[node].End = end;
            agentTree[node].MinX = agentTree[node].MaxX = agents[begin].Position.x;
            agentTree[node].MinY = agentTree[node].MaxY = agents[begin].Position.y;

            for (int i = begin + 1; i < end; ++i)
            {
                agentTree[node].MaxX = Math.Max(agentTree[node].MaxX, agents[i].Position.x);
                agentTree[node].MinX = Math.Min(agentTree[node].MinX, agents[i].Position.x);
                agentTree[node].MaxY = Math.Max(agentTree[node].MaxY, agents[i].Position.y);
                agentTree[node].MinY = Math.Min(agentTree[node].MinY, agents[i].Position.y);
            }

            if (end - begin > MAX_LEAF_SIZE)
            {
                /* No leaf node. */
                bool isVertical = agentTree[node].MaxX - agentTree[node].MinX > agentTree[node].MaxY - agentTree[node].MinY;
                float splitValue = 0.5f * (isVertical ? agentTree[node].MaxX + agentTree[node].MinX : agentTree[node].MaxY + agentTree[node].MinY);

                int left = begin;
                int right = end;

                while (left < right)
                {
                    while (left < right && (isVertical ? agents[left].Position.x : agents[left].Position.y) < splitValue)
                    {
                        ++left;
                    }

                    while (right > left && (isVertical ? agents[right - 1].Position.x : agents[right - 1].Position.y) >= splitValue)
                    {
                        --right;
                    }

                    if (left < right)
                    {
                        (agents[left], agents[right - 1]) = (agents[right - 1], agents[left]);
                        ++left;
                        --right;
                    }
                }

                int leftSize = left - begin;

                if (leftSize == 0)
                {
                    ++leftSize;
                    ++left;
                    ++right;
                }

                agentTree[node].Left = node + 1;
                agentTree[node].Right = node + 2 * leftSize;

                BuildAgentTreeRecursive(begin, left, agentTree[node].Left);
                BuildAgentTreeRecursive(left, end, agentTree[node].Right);
            }
        }

        /**
         * <summary>Recursive method for building an obstacle k-D tree.
         * </summary>
         *
         * <returns>An obstacle k-D tree node.</returns>
         *
         * <param name="obstacles">A list of obstacles.</param>
         */
        private ObstacleTreeNode BuildObstacleTreeRecursive(IList<Obstacle> obstacles)
        {
            if (obstacles.Count == 0)
            {
                return null;
            }

            ObstacleTreeNode node = new ObstacleTreeNode();

            int optimalSplit = 0;
            int minLeft = obstacles.Count;
            int minRight = obstacles.Count;

            for (int i = 0; i < obstacles.Count; ++i)
            {
                int leftSize = 0;
                int rightSize = 0;

                Obstacle obstacleI1 = obstacles[i];
                Obstacle obstacleI2 = obstacleI1.Next;

                /* Compute optimal split node. */
                for (int j = 0; j < obstacles.Count; ++j)
                {
                    if (i == j)
                    {
                        continue;
                    }

                    Obstacle obstacleJ1 = obstacles[j];
                    Obstacle obstacleJ2 = obstacleJ1.Next;

                    float j1LeftOfI = RVOMath.LeftOf(obstacleI1.Pos, obstacleI2.Pos, obstacleJ1.Pos);
                    float j2LeftOfI = RVOMath.LeftOf(obstacleI1.Pos, obstacleI2.Pos, obstacleJ2.Pos);

                    if (j1LeftOfI >= -RVOMath.kEpsilon && j2LeftOfI >= -RVOMath.kEpsilon)
                    {
                        ++leftSize;
                    }
                    else if (j1LeftOfI <= RVOMath.kEpsilon && j2LeftOfI <= RVOMath.kEpsilon)
                    {
                        ++rightSize;
                    }
                    else
                    {
                        ++leftSize;
                        ++rightSize;
                    }

                    if (new FloatPair(Math.Max(leftSize, rightSize), Math.Min(leftSize, rightSize)) >= new FloatPair(Math.Max(minLeft, minRight), Math.Min(minLeft, minRight)))
                    {
                        break;
                    }
                }

                if (new FloatPair(Math.Max(leftSize, rightSize), Math.Min(leftSize, rightSize)) < new FloatPair(Math.Max(minLeft, minRight), Math.Min(minLeft, minRight)))
                {
                    minLeft = leftSize;
                    minRight = rightSize;
                    optimalSplit = i;
                }
            }

            {
                /* Build split node. */
                IList<Obstacle> leftObstacles = new List<Obstacle>(minLeft);

                for (int n = 0; n < minLeft; ++n)
                {
                    leftObstacles.Add(null);
                }

                IList<Obstacle> rightObstacles = new List<Obstacle>(minRight);

                for (int n = 0; n < minRight; ++n)
                {
                    rightObstacles.Add(null);
                }

                int left_counter = 0;
                int right_counter = 0;
                int i = optimalSplit;

                Obstacle obstacle_i1 = obstacles[i];
                Obstacle obstacle_i2 = obstacle_i1.Next;

                for (int j = 0; j < obstacles.Count; ++j)
                {
                    if (i == j)
                    {
                        continue;
                    }

                    var obstacle_j1 = obstacles[j];
                    var obstacle_j2 = obstacle_j1.Next;

                    var j1_left_of_i = RVOMath.LeftOf(obstacle_i1.Pos, obstacle_i2.Pos, obstacle_j1.Pos);
                    var j2_left_of_i = RVOMath.LeftOf(obstacle_i1.Pos, obstacle_i2.Pos, obstacle_j2.Pos);

                    if (j1_left_of_i >= -RVOMath.kEpsilon && j2_left_of_i >= -RVOMath.kEpsilon)
                    {
                        leftObstacles[left_counter++] = obstacles[j];
                    }
                    else if (j1_left_of_i <= RVOMath.kEpsilon && j2_left_of_i <= RVOMath.kEpsilon)
                    {
                        rightObstacles[right_counter++] = obstacles[j];
                    }
                    else
                    {
                        /* Split obstacle j. */
                        var t = RVOMath.Det(obstacle_i2.Pos - obstacle_i1.Pos, obstacle_j1.Pos - obstacle_i1.Pos) / RVOMath.Det(obstacle_i2.Pos - obstacle_i1.Pos, obstacle_j1.Pos - obstacle_j2.Pos);

                        var split_point = obstacle_j1.Pos + t * (obstacle_j2.Pos - obstacle_j1.Pos);

                        var new_obstacle = new Obstacle
                        {
                            Pos = split_point,
                            Prev = obstacle_j1,
                            Next = obstacle_j2,
                            Convex = true,
                            Dir = obstacle_j1.Dir,
                            ID = sys.Obstacles.Count
                        };

                        sys.Obstacles.Add(new_obstacle);

                        obstacle_j1.Next = new_obstacle;
                        obstacle_j2.Prev = new_obstacle;

                        if (j1_left_of_i > 0.0f)
                        {
                            leftObstacles[left_counter++] = obstacle_j1;
                            rightObstacles[right_counter++] = new_obstacle;
                        }
                        else
                        {
                            rightObstacles[right_counter++] = obstacle_j1;
                            leftObstacles[left_counter++] = new_obstacle;
                        }
                    }
                }

                node.Obstacle = obstacle_i1;
                node.Left = BuildObstacleTreeRecursive(leftObstacles);
                node.Right = BuildObstacleTreeRecursive(rightObstacles);

                return node;
            }
        }

        private float GetSquareSum(AgentTreeNode atn, Vector2 position)
        {
            float sum = 0;

            sum += Math.Max(0.0f, atn.MinX - position.x).Square();
            sum += Math.Max(0.0f, position.x - atn.MaxX).Square();
            sum += Math.Max(0.0f, atn.MinY - position.y).Square();
            sum += Math.Max(0.0f, position.y - atn.MaxY).Square();

            return sum;
        }

        private void QueryAgentTreeRecursive(Vector2 position, ref float range_sq, ref int agent_no, int node)
        {
            var atn = agentTree[node];
            if (atn.End - atn.Begin <= MAX_LEAF_SIZE)
            {
                for (int i = atn.Begin; i < atn.End; ++i)
                {
                    var agent = agents[i];
                    float dist_sq = (position - agent.Position).sqrMagnitude;
                    if (dist_sq < range_sq)
                    {
                        range_sq = dist_sq;
                        agent_no = agent.Id;
                    }
                }
            }
            else
            {
                var atn_left = agentTree[atn.Left];
                var atn_right = agentTree[atn.Right];

                var dist_sq_left = GetSquareSum(atn_left, position);
                var dist_sq_right = GetSquareSum(atn_right, position);

                if (dist_sq_left < dist_sq_right)
                {
                    if (dist_sq_left < range_sq)
                    {
                        QueryAgentTreeRecursive(position, ref range_sq, ref agent_no, atn.Left);

                        if (dist_sq_right < range_sq)
                        {
                            QueryAgentTreeRecursive(position, ref range_sq, ref agent_no, atn.Right);
                        }
                    }
                }
                else
                {
                    if (dist_sq_right < range_sq)
                    {
                        QueryAgentTreeRecursive(position, ref range_sq, ref agent_no, atn.Right);

                        if (dist_sq_left < range_sq)
                        {
                            QueryAgentTreeRecursive(position, ref range_sq, ref agent_no, atn.Left);
                        }
                    }
                }

            }
        }

        /**
         * <summary>Recursive method for computing the agent neighbors of the
         * specified agent.</summary>
         *
         * <param name="agent">The agent for which agent neighbors are to be
         * computed.</param>
         * <param name="range_sq">The squared range around the agent.</param>
         * <param name="node">The current agent k-D tree node index.</param>
         */
        private void QueryAgentTreeRecursive(Agent agent, ref float range_sq, int node)
        {
            var atn = agentTree[node];
            if (atn.End - atn.Begin <= MAX_LEAF_SIZE)
            {
                for (int i = atn.Begin; i < atn.End; ++i)
                {
                    agent.InsertAgentNeighbor(agents[i], ref range_sq);
                }
            }
            else
            {
                var atn_left = agentTree[atn.Left];
                var atn_right = agentTree[atn.Right];

                float distSqLeft = GetSquareSum(atn_left, agent.Position);
                float distSqRight = GetSquareSum(atn_right, agent.Position);

                if (distSqLeft < distSqRight)
                {
                    if (distSqLeft < range_sq)
                    {
                        QueryAgentTreeRecursive(agent, ref range_sq, atn.Left);

                        if (distSqRight < range_sq)
                        {
                            QueryAgentTreeRecursive(agent, ref range_sq, atn.Right);
                        }
                    }
                }
                else
                {
                    if (distSqRight < range_sq)
                    {
                        QueryAgentTreeRecursive(agent, ref range_sq, atn.Right);

                        if (distSqLeft < range_sq)
                        {
                            QueryAgentTreeRecursive(agent, ref range_sq, atn.Left);
                        }
                    }
                }

            }
        }

        /**
         * <summary>Recursive method for computing the obstacle neighbors of the
         * specified agent.</summary>
         *
         * <param name="agent">The agent for which obstacle neighbors are to be
         * computed.</param>
         * <param name="range_sq">The squared range around the agent.</param>
         * <param name="node">The current obstacle k-D node.</param>
         */
        private void QueryObstacleTreeRecursive(Agent agent, float range_sq, ObstacleTreeNode node)
        {
            if (node != null)
            {
                Obstacle obstacle1 = node.Obstacle;
                Obstacle obstacle2 = obstacle1.Next;

                float agentLeftOfLine = RVOMath.LeftOf(obstacle1.Pos, obstacle2.Pos, agent.Position);

                QueryObstacleTreeRecursive(agent, range_sq, agentLeftOfLine >= 0.0f ? node.Left : node.Right);

                float distSqLine = agentLeftOfLine.Square() / (obstacle2.Pos - obstacle1.Pos).sqrMagnitude;

                if (distSqLine < range_sq)
                {
                    if (agentLeftOfLine < 0.0f)
                    {
                        /*
                         * Try obstacle at this node only if agent is on right side of
                         * obstacle (and can see obstacle).
                         */
                        agent.InsertObstacleNeighbor(node.Obstacle, range_sq);
                    }

                    /* Try other side of line. */
                    QueryObstacleTreeRecursive(agent, range_sq, agentLeftOfLine >= 0.0f ? node.Right : node.Left);
                }
            }
        }

        /**
         * <summary>Recursive method for querying the visibility between two
         * points within a specified radius.</summary>
         *
         * <returns>True if q1 and q2 are mutually visible within the radius;
         * false otherwise.</returns>
         *
         * <param name="q1">The first point between which visibility is to be
         * tested.</param>
         * <param name="q2">The second point between which visibility is to be
         * tested.</param>
         * <param name="radius">The radius within which visibility is to be
         * tested.</param>
         * <param name="node">The current obstacle k-D node.</param>
         */
        private bool QueryVisibilityRecursive(Vector2 q1, Vector2 q2, float radius, ObstacleTreeNode node)
        {
            if (node == null)
            {
                return true;
            }

            Obstacle obstacle1 = node.Obstacle;
            Obstacle obstacle2 = obstacle1.Next;

            float q1LeftOfI = RVOMath.LeftOf(obstacle1.Pos, obstacle2.Pos, q1);
            float q2LeftOfI = RVOMath.LeftOf(obstacle1.Pos, obstacle2.Pos, q2);
            float invLengthI = 1.0f / (obstacle2.Pos - obstacle1.Pos).sqrMagnitude;

            if (q1LeftOfI >= 0.0f && q2LeftOfI >= 0.0f)
            {
                return QueryVisibilityRecursive(q1, q2, radius, node.Left) &&
                    ((q1LeftOfI.Square() * invLengthI >= radius.Square() &&
                    q2LeftOfI.Square() * invLengthI >= radius.Square()) ||
                    QueryVisibilityRecursive(q1, q2, radius, node.Right));
            }

            if (q1LeftOfI <= 0.0f && q2LeftOfI <= 0.0f)
            {
                return QueryVisibilityRecursive(q1, q2, radius, node.Right) &&
                    ((q1LeftOfI.Square() * invLengthI >= radius.Square() &&
                    q2LeftOfI.Square() * invLengthI >= radius.Square()) ||
                    QueryVisibilityRecursive(q1, q2, radius, node.Left));
            }

            if (q1LeftOfI >= 0.0f && q2LeftOfI <= 0.0f)
            {
                /* One can see through obstacle from left to right. */
                return QueryVisibilityRecursive(q1, q2, radius, node.Left) &&
                    QueryVisibilityRecursive(q1, q2, radius, node.Right);
            }

            float point1LeftOfQ = RVOMath.LeftOf(q1, q2, obstacle1.Pos);
            float point2LeftOfQ = RVOMath.LeftOf(q1, q2, obstacle2.Pos);
            float invLengthQ = 1.0f / (q2 - q1).sqrMagnitude;

            return point1LeftOfQ * point2LeftOfQ >= 0.0f &&
                point1LeftOfQ.Square() * invLengthQ > radius.Square() &&
                point2LeftOfQ.Square() * invLengthQ > radius.Square() &&
                QueryVisibilityRecursive(q1, q2, radius, node.Left) &&
                QueryVisibilityRecursive(q1, q2, radius, node.Right);
        }
    }
}
